import asyncio
import os
import shlex
import shutil
import contextlib
from pathlib import Path
from typing import Literal, Optional

import psutil
from pydantic import BaseModel, Field

from mcp.server.fastmcp import Context, FastMCP

NAME = "MCP-openEuler-Server"

mcp = FastMCP(
    name=NAME,
    instructions=(
        "MCP-openEuler-Server exposes safe system administration tools for openEuler Linux. "
        "It supports command execution, file management, service control, resource monitoring, "
        "process management, and performance testing. Tools are sandboxed with optional allowlists "
        "and path constraints."
    ),
)


class ExecResult(BaseModel):
    returncode: int
    stdout: str
    stderr: str


class ExecSpec(BaseModel):
    cmd: str = Field(description="Command to execute. Use absolute paths for sensitive binaries.")
    cwd: Optional[str] = Field(default=None, description="Working directory")
    timeout_sec: int = Field(default=60, ge=1, le=600)
    shell: bool = Field(default=False, description="If True, run via shell. Defaults to False for safety.")


def _parse_list_env(name: str, default: list[str]) -> list[str]:
    val = os.getenv(name)
    if not val:
        return default
    # support colon or comma separated
    parts = [p.strip() for p in val.replace(";", ":").replace(",", ":").split(":")]
    return [p for p in parts if p]

ALLOWED_BIN_DIRS = _parse_list_env(
    "ALLOWED_BIN_DIRS",
    [
        "/usr/bin",
        "/usr/sbin",
        "/bin",
        "/sbin",
    ],
)

SAFE_PATHS = _parse_list_env("SAFE_PATHS", ["/home", "/tmp", "/var/tmp", "/etc", "/var/log", "/opt"])
ALLOW_SHELL = os.getenv("ALLOW_SHELL", "false").lower() in {"1", "true", "yes"}


def _is_under(path: Path, allowed_roots: list[str]) -> bool:
    try:
        p = path.resolve()
    except Exception:
        return False
    for root in allowed_roots:
        try:
            if p.is_relative_to(Path(root)):
                return True
        except AttributeError:
            # For Python <3.9 compatibility, emulate is_relative_to
            try:
                p.relative_to(Path(root))
                return True
            except Exception:
                pass
    return False


@mcp.tool(title="Execute Command", description="Execute a Linux command with timeout and capture output.")
async def exec_command(spec: ExecSpec, ctx: Context) -> ExecResult:
    await ctx.info(f"Executing: {spec.cmd}")
    if spec.shell and not ALLOW_SHELL:
        return ExecResult(returncode=126, stdout="", stderr="Shell execution disabled by policy")
    if not spec.shell:
        # If shell=False, resolve binary path and check allowlist dirs
        parts = shlex.split(spec.cmd)
        if not parts:
            return ExecResult(returncode=1, stdout="", stderr="Empty command")
        binary = parts[0]
        if os.path.isabs(binary):
            bin_path = Path(binary)
        else:
            # try which
            resolved = shutil.which(binary)
            bin_path = Path(resolved) if resolved else Path(binary)
        if not _is_under(bin_path, ALLOWED_BIN_DIRS):
            return ExecResult(returncode=126, stdout="", stderr=f"Executable not allowed: {bin_path}")

    try:
        if spec.shell:
            proc = await asyncio.create_subprocess_shell(
                spec.cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=spec.cwd,
            )
        else:
            parts = shlex.split(spec.cmd)
            proc = await asyncio.create_subprocess_exec(
                parts[0], *parts[1:],
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=spec.cwd,
            )
        try:
            stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=spec.timeout_sec)
        except asyncio.TimeoutError:
            proc.kill()
            return ExecResult(returncode=124, stdout="", stderr="Command timed out")
        return ExecResult(returncode=proc.returncode or 0, stdout=stdout.decode(), stderr=stderr.decode())
    except FileNotFoundError as e:
        return ExecResult(returncode=127, stdout="", stderr=str(e))


class FileWriteSpec(BaseModel):
    path: str
    content: str = ""
    overwrite: bool = True
    mkdirs: bool = True


@mcp.tool(title="Write File", description="Create or overwrite a text file. Constraints apply to allowed paths.")
async def write_file(spec: FileWriteSpec, ctx: Context) -> dict:
    p = Path(spec.path)
    if not _is_under(p, SAFE_PATHS):
        return {"ok": False, "error": "Path not allowed"}
    if p.exists() and not spec.overwrite:
        return {"ok": False, "error": "File exists"}
    if spec.mkdirs:
        p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(spec.content)
    await ctx.info(f"Wrote {p}")
    return {"ok": True, "path": str(p)}


class FileReadSpec(BaseModel):
    path: str
    max_bytes: int = 200_000


@mcp.tool(title="Read File", description="Read a text file.")
async def read_file(spec: FileReadSpec, ctx: Context) -> dict:
    p = Path(spec.path)
    if not p.exists():
        return {"ok": False, "error": "Not found"}
    if not _is_under(p, SAFE_PATHS):
        return {"ok": False, "error": "Path not allowed"}
    data = p.read_bytes()[: spec.max_bytes]
    return {"ok": True, "path": str(p), "content": data.decode(errors="replace")}


class FileDeleteSpec(BaseModel):
    path: str


@mcp.tool(title="Delete File", description="Delete a file.")
async def delete_file(spec: FileDeleteSpec, ctx: Context) -> dict:
    p = Path(spec.path)
    if not _is_under(p, SAFE_PATHS):
        return {"ok": False, "error": "Path not allowed"}
    if p.is_file():
        p.unlink(missing_ok=True)
        return {"ok": True}
    return {"ok": False, "error": "Not a file"}


class DirListSpec(BaseModel):
    path: str
    recursive: bool = False
    max_entries: int = 1000


@mcp.tool(title="List Directory", description="List directory contents.")
async def list_dir(spec: DirListSpec, ctx: Context) -> dict:
    p = Path(spec.path)
    if not _is_under(p, SAFE_PATHS):
        return {"ok": False, "error": "Path not allowed"}
    entries: list[dict] = []
    count = 0
    if spec.recursive:
        for root, dirs, files in os.walk(p):
            for name in dirs + files:
                fp = Path(root) / name
                st = fp.stat()
                entries.append({
                    "path": str(fp),
                    "is_dir": fp.is_dir(),
                    "size": st.st_size,
                    "mtime": st.st_mtime,
                })
                count += 1
                if count >= spec.max_entries:
                    break
            if count >= spec.max_entries:
                break
    else:
        for fp in p.iterdir():
            st = fp.stat()
            entries.append({
                "path": str(fp),
                "is_dir": fp.is_dir(),
                "size": st.st_size,
                "mtime": st.st_mtime,
            })
            count += 1
            if count >= spec.max_entries:
                break
    return {"ok": True, "entries": entries, "truncated": count >= spec.max_entries}


class ServiceSpec(BaseModel):
    name: str
    action: Literal["start", "stop", "restart", "status", "enable", "disable"]


@mcp.tool(title="Manage Service", description="Control systemd services via systemctl.")
async def manage_service(spec: ServiceSpec, ctx: Context) -> ExecResult:
    cmd = f"/usr/bin/systemctl {spec.action} {shlex.quote(spec.name)}"
    return await exec_command(ExecSpec(cmd=cmd, shell=True, timeout_sec=60), ctx)


@mcp.tool(title="System Stats", description="Get CPU, memory, disk, and network stats.")
async def system_stats(ctx: Context) -> dict:
    cpu = psutil.cpu_percent(interval=0.2)
    vm = psutil.virtual_memory()._asdict()
    disks = {p.mountpoint: psutil.disk_usage(p.mountpoint)._asdict() for p in psutil.disk_partitions(all=False)}
    net = psutil.net_io_counters(pernic=True)
    net_stats = {k: v._asdict() for k, v in net.items()}
    return {"cpu_percent": cpu, "memory": vm, "disks": disks, "net": net_stats}


class ProcSpec(BaseModel):
    pid: int
    signal: Literal["TERM", "KILL", "HUP"] = "TERM"


@mcp.tool(title="Manage Process", description="Send a signal to a process.")
async def manage_process(spec: ProcSpec, ctx: Context) -> dict:
    sig_map = {"TERM": 15, "KILL": 9, "HUP": 1}
    try:
        os.kill(spec.pid, sig_map[spec.signal])
        return {"ok": True}
    except ProcessLookupError:
        return {"ok": False, "error": "No such process"}
    except PermissionError:
        return {"ok": False, "error": "Permission denied"}


class BenchmarkSpec(BaseModel):
    type: Literal["cpu", "io"] = "cpu"
    seconds: int = Field(default=5, ge=1, le=120)


@mcp.tool(title="Benchmark", description="Run simple CPU or IO benchmark.")
async def benchmark(spec: BenchmarkSpec, ctx: Context) -> dict:
    if spec.type == "cpu":
        # Busy loop for specified seconds, counting iterations
        count = 0
        async def spin():
            nonlocal count
            while True:
                # simple arithmetic to keep busy
                count += 1
                if count % 1_000_000 == 0:
                    await asyncio.sleep(0)
        task = asyncio.create_task(spin())
        try:
            await asyncio.sleep(spec.seconds)
        finally:
            task.cancel()
            with contextlib.suppress(Exception):
                await task
        return {"ok": True, "iterations": count, "seconds": spec.seconds}
    else:
        # IO benchmark: write and read temp file
        import tempfile, time
        size = 32 * 1024 * 1024  # 32MB
        data = os.urandom(1024 * 1024)
        repeats = size // len(data)
        with tempfile.NamedTemporaryFile(delete=False) as f:
            start = time.time()
            for _ in range(repeats):
                f.write(data)
            f.flush()
            os.fsync(f.fileno())
            wdur = time.time() - start
            path = f.name
        start = time.time()
        with open(path, "rb") as f:
            while f.read(1024 * 1024):
                pass
        rdur = time.time() - start
        os.unlink(path)
        return {"ok": True, "write_sec": wdur, "read_sec": rdur, "size_bytes": size}


def main():
    # Default to stdio for agent use; allow env to switch to streamable-http
    transport = os.environ.get("MCP_TRANSPORT", "stdio")
    if transport == "streamable-http":
        mcp.run(transport="streamable-http")
    else:
        mcp.run()


if __name__ == "__main__":
    main()

